package cn.doitedu.ml.demo

import org.apache.spark.ml.evaluation.RegressionEvaluator
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.ml.regression.LinearRegression
import org.apache.spark.sql.SparkSession

import scala.collection.mutable

object LinerRegression {

  def main(args: Array[String]): Unit = {

    val spark = SparkSession.builder().appName("线性回归示例").master("local").getOrCreate()

    import spark.implicits._
    import org.apache.spark.sql.functions._

    val sample = spark.read.option("header", true).option("inferSchema", true).csv("user_portrait/data/linear/sample/sample.csv")
    sample.printSchema()
    sample.show(100,false)


    // 特征工程

    val arr2vec = udf((arr:mutable.WrappedArray[Double])=>{
      Vectors.dense(arr.toArray)
    })

    val vecDF = sample.select(arr2vec(array('area, 'floor)) as "features", 'price as "label")


    // 构造算法工具
    val linerRegression = new LinearRegression()
        .setFeaturesCol("features")
        .setLabelCol("label")
        .setRegParam(1.0)  //  防止过拟合
        .setMaxIter(100)


    val model = linerRegression.fit(vecDF)

    // 用模型来预测未知输出的数据

    val test = spark.read.option("header", true).option("inferSchema", true).csv("user_portrait/data/linear/test/test.csv")

    val testVec = test.select(arr2vec(array('area, 'floor)) as "features",'price as "label")

    val res = model.transform(testVec)
    res.cache()
    res.show(100,false)


    // 模型评估
    val evaluator = new RegressionEvaluator()
        .setLabelCol("label")
        .setPredictionCol("prediction")
        .setMetricName("rmse")

    val d: Double = evaluator.evaluate(res)
    println(d)

    spark.close()





  }

}
